# fl_server.py
from quantum_audit_system import QuantumAuditSystem
from noise_revocation_engine import NoiseRevocationEngine

class SecurityException(Exception):
    """Exception for security protocol violations"""
    pass

class FLServer:
    """
    Federated learning server with noise revocation
    
    Args:
        global_model (nn.Module): Global model instance
    """
    
    def __init__(self, global_model):
        self.global_model = global_model
        self.quantum_audit = QuantumAuditSystem()
    
    def aggregate_models(self, client_updates, quantum_tags):
        """
        Aggregate client updates and verify quantum consistency
        
        Args:
            client_updates (list): List of (noisy_params, quantum_tag)
            quantum_tags (list): List of quantum tags from clients
            
        Returns:
            list: Aggregated global model parameters
            
        Raises:
            SecurityException: If quantum consistency check fails
        """
        # Verify quantum consistency
        if not self.quantum_audit.verify_consistency(quantum_tags):
            raise SecurityException("Quantum audit consistency check failed")
        
        # Federated averaging
        avg_params = [torch.zeros_like(param) for param in client_updates[0]]
        for params, _ in client_updates:
            for i, param in enumerate(params):
                avg_params[i] += param / len(client_updates)
        
        return avg_params
    
    def revoke_global_noise(self, global_params, noise_vector):
        """
        Remove global noise from aggregated model
        
        Args:
            global_params (list): Noisy global parameters
            noise_vector (torch.Tensor): Isomorphic noise vector
            
        Returns:
            list: Clean global parameters
        """
        return NoiseRevocationEngine.remove_noise(global_params, noise_vector)
